/******************************************************************************
 * Decryptor class decrypts messages encrypted using the Encryptor class. It
 * uses OpenSSL enveloping to a RSA-decrypt one-time symmetric session key. The
 * session key is used to decrypt the payload. This work was done for Zero Leak
 * Architecture implementation for Dr. Hiroshi Fujinoki.
 *
 * Usage: Instantiated Decryptor must make the following calls in order:
 * keySrcFile(...), decryptSessionKey(...), decrypt(...).
 *****************************************************************************/
#include "Decryptor.h"
#include "netinet/in.h"

Decryptor::Decryptor(const char* rsa_keyfile, const EVP_CIPHER* cipher) :
	type(cipher) {

	ctx = new EVP_CIPHER_CTX;
	keySrcFile(rsa_keyfile);
}

/******************************************************************************
 * KeySrcFile: Locates rsa private key file and allocates memory for key-
 * 			related attributes.
 * Pre:		None.
 * Post:	Ready to call decryptSessionKey(...).
 *****************************************************************************/
Decryptor& Decryptor::keySrcFile(const char* rsa_keyfile) {

	privk_file = rsa_keyfile;
	BIO* rsa_bio = BIO_new_file(privk_file, "r");
	rsa_key = RSA_new();
	PEM_read_bio_RSAPrivateKey(rsa_bio, &rsa_key, NULL, NULL);
	BIO_free(rsa_bio);
	privk = EVP_PKEY_new();
	EVP_PKEY_assign_RSA(privk, rsa_key);

	return *this;
}

/******************************************************************************
 * DecryptSessionKey: Prepares keys for final decryption.
 * Pre:		Must call keySrcFile(...) first.
 * Post:	Ready to call decrypt(...).
 *****************************************************************************/
Decryptor& Decryptor::setupSession(unsigned char* init_vector, unsigned char* session_key, int session_key_len) {

	ek = session_key;
	ekl = session_key_len;

	iv = new unsigned char[EVP_MAX_IV_LENGTH];
	memcpy(iv, init_vector, EVP_MAX_IV_LENGTH);

	return *this;
}

unsigned char* Decryptor::DecryptPacket(unsigned char* pkt, int pkt_len, int* decrypt_len) {
   //get iv, encrypted key length, and encrypted key
	int encrypt_key_len, encrypt_data_len;
	unsigned char* encrypt_key;
	unsigned char* encrypt_data;

	//iv
	unsigned char* init_vect = pkt;
	int index = EVP_MAX_IV_LENGTH;

	//ekl
	memcpy(&encrypt_key_len, &pkt[index], sizeof(uint32_t));
	encrypt_key_len = ntohl(encrypt_key_len);
	index += sizeof(uint32_t);

	//ek
	encrypt_key = &pkt[index];
	index += encrypt_key_len;

	//encrypted data size
	memcpy(&encrypt_data_len, &pkt[index], sizeof(uint32_t));
	encrypt_data_len = ntohl(encrypt_data_len);
	index += sizeof(uint32_t);

	encrypt_data = &pkt[index];

	setupSession(init_vect, encrypt_key, encrypt_key_len);

	return decrypt(encrypt_data, encrypt_data_len, decrypt_len);
}
/******************************************************************************
 * Decrypt: Performs decryption on message.
 * Pre: 	Must call decryptSessionKey(...) first.
 * Post:	Null-terminated decrypted message is stored in out user-supplied
 * 			buffer.
 *****************************************************************************/
unsigned char* Decryptor::decrypt(unsigned char* encrypt_buf,	int encrypt_len, int* decrypt_len) {
  int bytes_written;
	unsigned char* decrypt_buf = new unsigned char[encrypt_len];

	EVP_OpenInit(ctx, type, ek, ekl, iv, privk);
	EVP_OpenUpdate(ctx, decrypt_buf, &bytes_written, encrypt_buf, encrypt_len);
	*decrypt_len = bytes_written;

	// EVP_OpenFinal changes overhead value to # bytes of encryption overhead ----
	EVP_OpenFinal(ctx, &decrypt_buf[*decrypt_len], &bytes_written);
	*decrypt_len += bytes_written;

	//this is ok because we are sure that encrypt_len will always be greater than decrypt_len
	decrypt_buf[*decrypt_len] = '\0';

	printf("decrypt::decrypt_len: %d\n", *decrypt_len);

	return (decrypt_buf);
}

Decryptor::~Decryptor() {
	delete ctx;
	delete type;
	delete[] iv;
	delete[] privk_file;
	delete ek;
	EVP_PKEY_free(privk);
	RSA_free(rsa_key);
}
